/*
 * FileName : AES.h
 * 
 */
#ifndef __AES_H__
#define __AES_H__

#include <cstdio>

extern void print(unsigned char* state, int len);

class AES  
{
public:
	AES(unsigned char* key = NULL);
	virtual ~AES();
	void SetKey(unsigned char *key);
	unsigned char* Cipher(unsigned char* input, unsigned char* output);
	unsigned char* InvCipher(unsigned char* input, unsigned char* output);
	void* Cipher(void* input, void *output, int length=0);
	void* InvCipher(void* input,void *output, int length);

private:
	unsigned char Sbox[256];
	unsigned char InvSbox[256];
	unsigned char w[11][4][4];

	void KeyExpansion(unsigned char* key, unsigned char w[][4][4]);
	unsigned char FFmul(unsigned char a, unsigned char b);

	void SubBytes(unsigned char state[][4]);
	void ShiftRows(unsigned char state[][4]);
	void MixColumns(unsigned char state[][4]);
	void AddRoundKey(unsigned char state[][4], unsigned char k[][4]);

	void InvSubBytes(unsigned char state[][4]);
	void InvShiftRows(unsigned char state[][4]);
	void InvMixColumns(unsigned char state[][4]);
};

enum AESMode_t { MODE_OFB = 1, MODE_CFB, MODE_CBC, MODE_ECB };
class AESModeOfOperation {
private:
	AES *m_aes;
	AESMode_t	  m_mode;
	unsigned char m_key[16];
	unsigned char m_iv[16];
//	bool		  m_firstround;
public:
	AESModeOfOperation();
	~AESModeOfOperation();
	void set_mode(AESMode_t _mode);
	//AESMode_t get_mode();
	void set_key (unsigned char *key);
	void set_iv(unsigned char *iv);
	int  Encrypt(unsigned char *input, int length, unsigned char *output);
	int  Decrypt(unsigned char *input, int length, unsigned char *output);
};


#endif 